package com.zhuangjie.openfeign.aop.retry.pool;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.function.Consumer;

/**
 * 可重试线程池
 */
@Slf4j
public class RetryableThreadPool {

    private final ExecutorService executor;  // 线程池

    /**
     * 构造函数
     * @param executor 线程池执行器
     */
    public RetryableThreadPool(ExecutorService executor) {
        this.executor = executor;
    }

    /**
     * 执行可运行任务，如果遇到异常进行重试
     *
     * @param runnable 可运行的任务
     * @param clazz 指定的异常类
     * @param retryTimes 重试次数
     * @param retryIntervalMillis 重试间隔时间
     * @param onFailure 不可恢复错误处理的回调
     */
    public <T extends Exception> void execute(Runnable runnable, Class<T> clazz, int retryTimes, int retryIntervalMillis, Consumer<Exception> onFailure) {
        executor.submit(() -> {
            retry(() -> {
                runnable.run();
                return null;
            }, new Class[]{clazz}, retryTimes, retryIntervalMillis, onFailure);
        });
    }

    /**
     * 执行C callable任务，如果遇到异常进行重试
     *
     * @param callable callable的任务
     * @param retryTimes 重新尝试的次数
     * @param retryIntervalMillis 重试的间隔时间
     * @param onFailure 不可以恢复错误处理的回调
     * @return 返回call的结果
     */
    public <R> R execute(Callable<R> callable,  Class<? extends Exception>[] exceptions, int retryTimes, int retryIntervalMillis, Consumer<Exception> onFailure) {
        try {
            return executor.submit(() -> retry(callable, exceptions, retryTimes, retryIntervalMillis, onFailure)).get();
        } catch (Exception e) {
            onFailure.accept(e);
            return null;
        }
    }

    /**
     * 重试逻辑
     *
     * @param callable 可调用的任务
     * @param retryTimes 重新尝试的次数
     * @param retryIntervalMillis 重试的间隔时间
     * @param onFailure 不可以恢复错误处理的回调
     * @return 返回call的结果
     */
    private <T extends Exception, R> R retry(Callable<R> callable, Class<? extends Exception>[] exceptions, int retryTimes, int retryIntervalMillis, Consumer<Exception> onFailure) {
        for (int i = 0; i <= retryTimes; i++) {
            try {
                return callable.call();  // if run success, return.
            } catch (Exception e) {
                //如果不是指定的异常类型，则直接执行失败回调并返回
                if (!isInstanceOfSpecifiedException(e, exceptions)) {
                    onFailure.accept(e);
                    return null;
                }
                //如果已达到最大重试次数，则执行失败回调并返回
                if (i >= retryTimes) {
                    log.info("Retry times exceed limit, give up.");
                    onFailure.accept(e);
                    return null;
                }
                log.info("Task failed, retry {}", i);
                try {
                    Thread.sleep(retryIntervalMillis);
                } catch (InterruptedException ex) {
                    throw new RuntimeException(ex);
                }
            }
        }
        return null;
    }

    /**
     * 判断异常是否是指定的异常类型
     *
     * @param sourceException 源异常
     * @param exceptions 目标异常类
     * @return 如果源异常是目标异常类型或目标异常类型的子类, 返回true，否则返回false
     */
    private <T extends Exception> boolean isInstanceOfSpecifiedException(Throwable sourceException, Class<? extends Exception>[] exceptions) {
        Throwable currentException = sourceException;
        while (currentException != null) {
            for (Class<? extends Exception> exception : exceptions) {
                if (exception.isInstance(currentException)) {
                    return true;
                }
            }
            currentException = currentException.getCause();
        }
        return false;
    }
}
